

import os


import cv2
import re
import numpy as np
import json

from glob import glob
from PIL import Image
from tqdm import tqdm
import glob


from .dataset import Dataset


class Video(object):

    def __init__(self, name, video_dir):
        self.name = name
        self.video_dir = video_dir
        self.gts = self.load_video()
        print(self.gts[0],len(self.gts))

    def load_video(self):
        path = os.path.join(self.video_dir, self.name)
        self.img_paths = sorted(glob.glob(os.path.join(path, "*.jpg")),
                              key=lambda x: int(os.path.basename(x).split('.')[0]))
        gt_path = os.path.join(self.video_dir, self.name, "groundtruth.txt")
        gts = []
        with open(gt_path, 'r') as f:
            while True:
                line = f.readline()
                if line == '':
                    gts = np.array(gts, dtype=np.float32)
                    return gts
                if ',' in line:
                    gt_pos = line.split(',')
                else:
                    gt_pos = line.split()
                gt_pos_int = [(float(element)) for element in gt_pos]

                # x, y = np.array(gt_pos_int[0:2])
                # w, h = np.array(gt_pos_int[4:6])-np.array(gt_pos_int[0:2])

                gts.append(gt_pos_int)

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img = cv2.imread(self.img_paths[idx])
        if len(self.gts)==1:
            gt_traj = self.gts[0]
        else:
            gt_traj = self.gts[idx]
        return img, gt_traj




class OTBDataset(Dataset):
    """
    Args:
        name: train or test
        dataset_root: dataset root
        load_img: wether to load all imgs
    """
    def __init__(self, name, dataset_root, load_img=True):
        super(OTBDataset, self).__init__(name, dataset_root)
        if name == "train":
            datapath = os.path.join(dataset_root,"trainval","trainval")

        elif name == "test":
            datapath = os.path.join(dataset_root,"test_public", "test_public")
        else:
            raise "name is error"

        video_names = os.listdir(datapath)

        # load videos
        self.videos = {}
        for video in video_names:
            print(video)
            self.videos[video] = Video(video,
                                      datapath)


        # set attr
        # attr = []
        # for x in self.videos.values():
        #     attr += x.attr
        # attr = set(attr)
        # self.attr = {}
        # self.attr['ALL'] = list(self.videos.keys())
        # for x in attr:
        #     self.attr[x] = []
        # for k, v in self.videos.items():
        #     for attr_ in v.attr:
        #         self.attr[attr_].append(k)

class DatasetFactory(object):
    @staticmethod
    def create_dataset(**kwargs):
        dataset = OTBDataset(**kwargs)
        return dataset

if __name__ == '__main__':
    dataset = DatasetFactory.create_dataset(name="test", dataset_root="F:\\6-single-object-tracking")